perf: optimize array_remove for scalar needle by lyne7-sc · Pull Request #22390 · apache/datafusion

lyne7-sc · 2026-05-20T09:29:00Z

Which issue does this PR close?

Closes #.

Rationale for this change

Similar to #22387 (array_replace scalar optimization)

array_remove / array_remove_n / array_remove_all perform element-wise comparison by invoking compare_element_to_list against each row's sub-array individually. When the needle is a scalar, this can be optimized by performing a single vectorized distinct comparison over the entire flattened values buffer.

What changes are included in this PR?

Add a specialized removal kernel (general_remove_with_scalar) that uses arrow_ord::cmp::distinct with Scalar wrapper for a single bulk comparison pass over the flat values buffer.
Extend SLT tests with multi-row scalar-argument coverage, NULL-containing arrays, empty-array edge cases, boundary n values, and LargeList type coverage.

Benchmarks

group                                                                    main                                   optimized
-----                                                                    ----                                   ---------
array_remove_all_int64/remove/list size: 10, num_rows: 4000              4.35   856.8±97.81µs        ? ?/sec    1.00    196.9±4.48µs        ? ?/sec
array_remove_all_int64/remove/list size: 100, num_rows: 10000            1.90      5.5±0.09ms        ? ?/sec    1.00      2.9±0.09ms        ? ?/sec
array_remove_all_int64/remove/list size: 500, num_rows: 10000            1.35     19.2±0.21ms        ? ?/sec    1.00     14.2±0.48ms        ? ?/sec
array_remove_all_int64_nested/remove/list size: 10, num_rows: 4000       1.00      7.1±0.12ms        ? ?/sec    1.04      7.4±0.12ms        ? ?/sec
array_remove_all_int64_nested/remove/list size: 100, num_rows: 3000      1.00     36.5±0.39ms        ? ?/sec    1.05     38.3±2.61ms        ? ?/sec
array_remove_all_int64_nested/remove/list size: 300, num_rows: 1500      1.01     53.5±2.26ms        ? ?/sec    1.00     53.0±0.99ms        ? ?/sec
array_remove_boolean/remove/list size: 10, num_rows: 4000                3.83    813.9±7.08µs        ? ?/sec    1.00    212.4±2.28µs        ? ?/sec
array_remove_boolean/remove/list size: 100, num_rows: 10000              2.73      3.7±0.03ms        ? ?/sec    1.00  1364.7±177.83µs        ? ?/sec
array_remove_boolean/remove/list size: 500, num_rows: 10000              2.34      9.8±0.14ms        ? ?/sec    1.00      4.2±0.25ms        ? ?/sec
array_remove_fixed_size_binary/remove/list size: 10, num_rows: 4000      3.16   918.2±16.76µs        ? ?/sec    1.00    290.6±9.79µs        ? ?/sec
array_remove_fixed_size_binary/remove/list size: 100, num_rows: 10000    1.56      6.9±0.13ms        ? ?/sec    1.00      4.4±0.15ms        ? ?/sec
array_remove_fixed_size_binary/remove/list size: 500, num_rows: 10000    1.17     27.7±0.84ms        ? ?/sec    1.00     23.6±2.04ms        ? ?/sec
array_remove_int64/remove/list size: 10, num_rows: 4000                  4.55    825.7±6.30µs        ? ?/sec    1.00    181.3±4.32µs        ? ?/sec
array_remove_int64/remove/list size: 100, num_rows: 10000                3.35      3.8±0.11ms        ? ?/sec    1.00  1135.6±54.87µs        ? ?/sec
array_remove_int64/remove/list size: 500, num_rows: 10000                2.04     10.3±0.35ms        ? ?/sec    1.00      5.1±0.39ms        ? ?/sec
array_remove_int64_nested/remove/list size: 10, num_rows: 4000           1.00      7.1±0.18ms        ? ?/sec    1.02      7.2±0.07ms        ? ?/sec
array_remove_int64_nested/remove/list size: 100, num_rows: 3000          1.00     36.1±1.35ms        ? ?/sec    1.07     38.5±3.67ms        ? ?/sec
array_remove_int64_nested/remove/list size: 300, num_rows: 1500          1.00     51.7±0.57ms        ? ?/sec    1.05     54.1±2.13ms        ? ?/sec
array_remove_n_int64/remove/list size: 10, num_rows: 4000                4.43    845.3±5.00µs        ? ?/sec    1.00    190.6±2.84µs        ? ?/sec
array_remove_n_int64/remove/list size: 100, num_rows: 10000              2.29      4.7±0.11ms        ? ?/sec    1.00      2.0±0.12ms        ? ?/sec
array_remove_n_int64/remove/list size: 500, num_rows: 10000              1.63     14.8±0.42ms        ? ?/sec    1.00      9.0±0.51ms        ? ?/sec
array_remove_n_int64_nested/remove/list size: 10, num_rows: 4000         1.00      7.0±0.09ms        ? ?/sec    1.29      8.9±3.44ms        ? ?/sec
array_remove_n_int64_nested/remove/list size: 100, num_rows: 3000        1.00     36.6±0.42ms        ? ?/sec    1.03     37.7±0.68ms        ? ?/sec
array_remove_n_int64_nested/remove/list size: 300, num_rows: 1500        1.00     52.7±3.68ms        ? ?/sec    1.03     54.5±4.49ms        ? ?/sec
array_remove_strings/remove/list size: 10, num_rows: 4000                2.50  1144.6±21.95µs        ? ?/sec    1.00   457.0±14.15µs        ? ?/sec
array_remove_strings/remove/list size: 100, num_rows: 10000              1.42     10.5±1.16ms        ? ?/sec    1.00      7.4±0.34ms        ? ?/sec
array_remove_strings/remove/list size: 500, num_rows: 10000              1.12     39.8±0.91ms        ? ?/sec    1.00     35.5±1.51ms        ? ?/sec

Are these changes tested?

Yes, existing and new SLT edge-case tests in array_remove.slt.

Are there any user-facing changes?

No.

neilconway

Nice performance win!

neilconway · 2026-05-20T16:09:11Z

+            );
+        }
+    };
+    let original_data = list_array.values().to_data();


This will be inefficient for sliced arrays.

I now slice the values to the range actually referenced by the offsets.

That said, I wanted to understand your concern better: when a GenericListArray is sliced, values() returns the full underlying array, and to_data() on it wraps the existing buffer references into ArrayData without copying. So the main downside I could identify is that Capacities::Array(original_data.len()) over-estimates the pre-allocation for sliced inputs. Were you thinking of a different inefficiency, or is the over-allocation what you had in mind?

The over-allocation was one part, but the bigger concern is calling the distinct kernel on the entire values buffer (see other comment).

neilconway · 2026-05-20T16:12:41Z

+            let list_array = array.as_list::<i64>();
+            general_remove_with_scalar::<i64>(list_array, needle, arr_n)
+        }
+        array_type => exec_err!("array_remove does not support type '{array_type}'."),


This is called by more than just array_remove; can we improve the error message?

sure, updated.

neilconway · 2026-05-20T16:23:46Z

+        for (i, keep) in eq_array.iter().enumerate() {
+            if keep == Some(false) && removed < max_removals {
+                if let Some(bs) = pending_batch_to_retain {
+                    mutable.extend(0, start + bs, start + i);
+                    copied += i - bs;
+                    pending_batch_to_retain = None;
+                }
+                removed += 1;
+            } else if pending_batch_to_retain.is_none() {
+                pending_batch_to_retain = Some(i);
+            }
+        }


I wonder if it would be possible to iterate only over the "false" bits, e.g., by negating the buffer and looking at BooleanBuffer::set_indices.

Great suggestion. Benchmarks show a ~20–40% improvement with this optimization.

neilconway · 2026-05-20T16:24:25Z

+        let mut copied = 0usize;
+        let mut pending_batch_to_retain: Option<usize> = None;
+        for (i, keep) in eq_array.iter().enumerate() {
+            if keep == Some(false) && removed < max_removals {


Can we break from the loop once we hit max_removals?

Good point. now break early once max_removals is reached.

neilconway · 2026-05-21T15:46:00Z

+        // Iterate only over the positions that need removal using set_indices,
+        // which is more efficient than scanning every bit.


Might be worth elaborating that the win here is mostly because we expect the # of values-to-remove is a lot smaller than the total array size, which it usually (but not always) will be.

neilconway · 2026-05-21T15:47:43Z

+    let keep_mask =
+        arrow_ord::cmp::distinct(list_array.values(), &Scalar::new(Arc::clone(needle)))?;


This will call the distinct kernel on all the elements in the value buffer, not just the ones that are visible in a sliced array.

neilconway · 2026-05-21T15:50:24Z

+            );
+        }
+    };
+    let original_data = list_array.values().to_data();


The over-allocation was one part, but the bigger concern is calling the distinct kernel on the entire values buffer (see other comment).

neilconway

Awesome, nice work on this! lgtm. Maybe update the benchmark numbers in the PR description if you get a chance.

lyne7-sc · 2026-05-22T15:32:31Z

Thanks @neilconway! Really appreciate the thorough and speedy review 🙏 Great suggestion — running benchmarks now and will update the pr description with fresh numbers shortly.

Jefffrey · 2026-05-23T05:47:20Z

+///
+/// Uses a bulk `distinct` comparison for non-null, non-nested scalar elements,
+/// falling back to the per-row `general_remove` path for null or nested types.
+fn remove_with_scalar_needle(


I find this function a little confusing, since the code paths can then be like:

invoke with args -> remove_with_scalar_needle -> array_remove_internal (fallback)

Perhaps we can hoist this null/nested check earlier to remove need for this function? Or the benefits of centralizing the check outweighs it?

Jefffrey · 2026-05-23T05:47:44Z

+        let num_rows = args.number_rows;
+        let list_array = list_arg.to_array(num_rows)?;
+        let max_array = max_arg.to_array(num_rows)?;
+        let arr_n = as_int64_array(&max_array)?.values().to_vec();


We're ignoring nulls in max_arg here, though I think this may be an existing issue

Jefffrey · 2026-05-23T05:49:58Z

+    let values_range_len = last_offset - first_offset;
+    let values_slice = list_array.values().slice(first_offset, values_range_len);
+    let original_data = values_slice.to_data();
+    let mut offsets = Vec::<OffsetSize>::with_capacity(list_array.len() + 1);


Use OffsetBufferBuilder here

Jefffrey · 2026-05-23T05:52:13Z

+    let mut offsets = Vec::<OffsetSize>::with_capacity(list_array.len() + 1);
+    offsets.push(OffsetSize::zero());
+
+    let mut mutable = MutableArrayData::with_capacities(


I wonder if an approach using take kernel could provide even more performance gains?

Jefffrey · 2026-05-23T05:57:07Z

+    let nulls = list_array.nulls().cloned();
+    let keep_mask =
+        arrow_ord::cmp::distinct(&values_slice, &Scalar::new(Arc::clone(needle)))?;
+    let remove_bits = match keep_mask.nulls() {


distinct is similar to neq, only differing in null handling. In particular, two
operands are considered DISTINCT if they have a different value or if one of them is NULL
and the other isn’t. The result of distinct is never NULL.

https://docs.rs/arrow/latest/arrow/compute/kernels/cmp/fn.distinct.html

I don't think we should have handling around the null buffer of keep_mask according to the documentation

Jefffrey · 2026-05-23T05:57:45Z

+    needle: &ArrayRef,
+    arr_n: &[i64],


Similar question to that of array_replace PR, in if we should just handle when needle & max are scalars only

lyne7-sc and others added 4 commits May 20, 2026 11:16

Refactor array remove invocation

9b3f32f

Merge branch 'apache:main' into perf/remove

0c787f4

enhance array_remove slt

9697c3a

Merge branch 'main' into perf/remove

a6e6ad1

github-actions Bot added sqllogictest SQL Logic Tests (.slt) functions Changes to functions implementation labels May 20, 2026

neilconway reviewed May 20, 2026

View reviewed changes

lyne7-sc added 2 commits May 21, 2026 22:41

apply suggestions

b6023a0

apply suggestions

2e7cd40

neilconway reviewed May 21, 2026

View reviewed changes

redundant comparisons

cdb6021

neilconway approved these changes May 22, 2026

View reviewed changes

lyne7-sc added a commit to lyne7-sc/datafusion that referenced this pull request May 23, 2026

apply optimization suggestion from apache#22390

a4df12f

Jefffrey reviewed May 23, 2026

View reviewed changes

		// Iterate only over the positions that need removal using set_indices,
		// which is more efficient than scanning every bit.

		let keep_mask =
		arrow_ord::cmp::distinct(list_array.values(), &Scalar::new(Arc::clone(needle)))?;

Conversation

lyne7-sc commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Benchmarks

Are these changes tested?

Are there any user-facing changes?

Uh oh!

neilconway left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

neilconway left a comment

Choose a reason for hiding this comment

Uh oh!

lyne7-sc commented May 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lyne7-sc commented May 20, 2026 •

edited

Loading